use async_trait::async_trait;
use chrono::{DateTime, TimeZone, Utc};
use pyo3::{Python, PyResult, types::PyModule};
use std::collections::HashMap;
use std::error::Error;
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::config::data_source_config::DataSourceConfig;
use super::{DataSource, StockQuote, SectorIndex, MarketSentiment, VolumeDistribution, MarketBreadth, VolatilityIndex};

#[derive(Clone)]
pub struct AkShareDataSource {
    config: DataSourceConfig,
    cache: Arc<Mutex<HashMap<String, (DateTime<Utc>, serde_json::Value)>>>,
}

impl AkShareDataSource {
    pub fn new(config: &DataSourceConfig) -> Self {
        Self {
            config: config.clone(),
            cache: Arc::new(Mutex::new(HashMap::new())),
        }
    }
    
    async fn fetch_with_cache<T, F>(&self, cache_key: &str, ttl_seconds: u64, fetch_fn: F) -> Result<T, Box<dyn Error + Send + Sync>>
    where
        T: for<'de> serde::Deserialize<'de> + Send + Sync + 'static,
        F: FnOnce() -> PyResult<T> + Send + 'static,
    {
        let now = Utc::now();
        
        // 检查缓存
        let mut cache = self.cache.lock().await;
        if let Some((timestamp, cached_data)) = cache.get(cache_key) {
            let age = now.signed_duration_since(*timestamp).num_seconds() as u64;
            if age < ttl_seconds {
                // 缓存有效，反序列化并返回
                return Ok(serde_json::from_value(cached_data.clone())?);
            }
        }
        
        // 在新线程中执行Python代码
        let result = tokio::task::spawn_blocking(move || {
            Python::with_gil(|py| {
                // 导入akshare模块
                let akshare = PyModule::import(py, "akshare")?;
                fetch_fn()
            })
        }).await?;
        
        let data = match result {
            Ok(data) => data,
            Err(e) => return Err(format!("Failed to fetch data: {}", e).into()),
        };
        
        // 更新缓存
        let json_data = serde_json::to_value(&data)?;
        cache.insert(cache_key.to_string(), (now, json_data.clone()));
        
        Ok(serde_json::from_value(json_data)?)
    }
    
    // 新增方法：调用Python端的板块指数数据采集
    async fn fetch_sector_indices(&self) -> Result<Vec<SectorIndex>, Box<dyn Error + Send + Sync>> {
        let ttl_seconds = 300; // 5分钟缓存
        
        self.fetch_with_cache("sector_indices", ttl_seconds, || {
            Python::with_gil(|py| {
                // 导入自定义的数据采集模块
                let data_collection = PyModule::import(py, "data_collection")?;
                
                // 调用获取板块指数的函数
                let indices_py = data_collection.getattr("get_all_sector_indices")?.call0(py)?;
                
                // 将Python对象转换为Rust对象
                let indices: Vec<SectorIndex> = indices_py.extract()?;
                
                Ok(indices)
            })
        }).await
    }
    
    // 新增方法：调用Python端的市场情绪指标采集
    async fn fetch_market_sentiment(&self) -> Result<MarketSentiment, Box<dyn Error + Send + Sync>> {
        let ttl_seconds = 600; // 10分钟缓存
        
        self.fetch_with_cache("market_sentiment", ttl_seconds, || {
            Python::with_gil(|py| {
                // 导入自定义的数据采集模块
                let data_collection = PyModule::import(py, "data_collection")?;
                
                // 调用获取市场情绪的函数
                let sentiment_py = data_collection.getattr("get_market_sentiment")?.call0(py)?;
                
                // 将Python对象转换为Rust对象
                let sentiment: MarketSentiment = sentiment_py.extract()?;
                
                Ok(sentiment)
            })
        }).await
    }
    
    // 新增方法：调用Python端的成交量分布数据采集
    async fn fetch_volume_distribution(&self) -> Result<VolumeDistribution, Box<dyn Error + Send + Sync>> {
        let ttl_seconds = 300; // 5分钟缓存
        
        self.fetch_with_cache("volume_distribution", ttl_seconds, || {
            Python::with_gil(|py| {
                // 导入自定义的数据采集模块
                let data_collection = PyModule::import(py, "data_collection")?;
                
                // 调用获取成交量分布的函数
                let distribution_py = data_collection.getattr("get_volume_distribution")?.call0(py)?;
                
                // 将Python对象转换为Rust对象
                let distribution: VolumeDistribution = distribution_py.extract()?;
                
                Ok(distribution)
            })
        }).await
    }
}

#[async_trait]
impl DataSource for AkShareDataSource {
    async fn get_stock_quote(&self, symbol: &str) -> Result<StockQuote, Box<dyn Error + Send + Sync>> {
        // 原有的股票报价获取实现...
        unimplemented!("This method is not implemented yet")
    }
    
    async fn get_sector_indices(&self) -> Result<Vec<SectorIndex>, Box<dyn Error + Send + Sync>> {
        self.fetch_sector_indices().await
    }
    
    async fn get_market_sentiment(&self) -> Result<MarketSentiment, Box<dyn Error + Send + Sync>> {
        self.fetch_market_sentiment().await
    }
    
    async fn get_volume_distribution(&self) -> Result<VolumeDistribution, Box<dyn Error + Send + Sync>> {
        self.fetch_volume_distribution().await
    }
    
    async fn get_market_breadth(&self) -> Result<MarketBreadth, Box<dyn Error + Send + Sync>> {
        // 原有的市场宽度获取实现...
        unimplemented!("This method is not implemented yet")
    }
    
    async fn get_volatility_index(&self, index_name: &str) -> Result<VolatilityIndex, Box<dyn Error + Send + Sync>> {
        // 原有的波动率指数获取实现...
        unimplemented!("This method is not implemented yet")
    }
}